import torch
import torch.nn as nn
import copy

class TransformerDecoderLayer(nn.Module):
  def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
    super(TransformerDecoderLayer, self).__init__()
    self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
    self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
    # Implementation of Feedforward model
    self.linear1 = nn.Linear(d_model, dim_feedforward)
    self.dropout = nn.Dropout(dropout)
    self.linear2 = nn.Linear(dim_feedforward, d_model)

    self.norm1 = nn.LayerNorm(d_model)
    self.norm2 = nn.LayerNorm(d_model)
    self.norm3 = nn.LayerNorm(d_model)
    self.dropout1 = nn.Dropout(dropout)
    self.dropout2 = nn.Dropout(dropout)
    self.dropout3 = nn.Dropout(dropout)

  def forward(self, tgt, memory, tgt_mask=None, memory_mask=None,
              tgt_key_padding_mask=None, memory_key_padding_mask=None):
    tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,
                          key_padding_mask=tgt_key_padding_mask)[0]
    tgt = tgt + self.dropout1(tgt2)
    tgt = self.norm1(tgt)
    tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,
                              key_padding_mask=memory_key_padding_mask)[0]
    tgt = tgt + self.dropout2(tgt2)
    tgt = self.norm2(tgt)
    tgt2 = self.linear2(self.dropout(F.relu(self.linear1(tgt))))
    tgt = tgt + self.dropout3(tgt2)
    tgt = self.norm3(tgt)
    return tgt

class TransformerDecoder(nn.Module):
  def __init__(self, decoder_layer, num_layers, norm=None):
    super(TransformerDecoder, self).__init__()
    self.layers = _get_clones(decoder_layer, num_layers)
    self.num_layers = num_layers
    self.norm = norm

  def forward(self, tgt, memory, tgt_mask=None,
              memory_mask=None, tgt_key_padding_mask=None,
              memory_key_padding_mask=None):
    output = tgt
    for mod in self.layers:
      output = mod(output, memory, tgt_mask=tgt_mask,
                   memory_mask=memory_mask,
                   tgt_key_padding_mask=tgt_key_padding_mask,
                   memory_key_padding_mask=memory_key_padding_mask)

    if self.norm is not None:
      output = self.norm(output)

    return output

def _get_clones(module, N):
  return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

# Example usage
decoder_layer = TransformerDecoderLayer(d_model=512, nhead=3)
decoder = TransformerDecoder(decoder_layer, num_layers=1)
memory = torch.rand(10, 32, 512)
tgt = torch.rand(20, 32, 512)
output = decoder(tgt, memory)

print(output.shape)  # torch.Size([20, 32, 512])